import pygame
import time
import random

# 初始化pygame
pygame.init()

# 常量定义
WIDTH, HEIGHT = 600, 600
GRID_SIZE = 20
CELL_SIZE = WIDTH // GRID_SIZE
WHITE = (255, 255, 255)
BLACK = (0, 0, 0)
RED = (255, 0, 0)
GREEN = (0, 255, 0)
BLUE = (0, 0, 255)

# 创建屏幕
screen = pygame.display.set_mode((WIDTH, HEIGHT))
pygame.display.set_caption("迷宫生成算法演示")

# 方向常量
UP = (0, -1)
DOWN = (0, 1)
LEFT = (-1, 0)
RIGHT = (1, 0)

class Maze:
    def __init__(self, size):
        self.size = size
        self.grid = [[{'walls': [True, True, True, True], 'visited': False} 
                      for _ in range(size)] for _ in range(size)]
        self.stack = []
        self.current_cell = None
        self.finished = False
        
    def reset(self):
        self.grid = [[{'walls': [True, True, True, True], 'visited': False} 
                      for _ in range(self.size)] for _ in range(self.size)]
        self.stack = []
        self.current_cell = None
        self.finished = False
    
    def get_neighbors(self, x, y):
        neighbors = []
        directions = [UP, DOWN, LEFT, RIGHT]
        
        for dx, dy in directions:
            nx, ny = x + dx, y + dy
            if 0 <= nx < self.size and 0 <= ny < self.size and not self.grid[ny][nx]['visited']:
                neighbors.append((nx, ny, dx, dy))
        
        return neighbors
    
    def remove_wall(self, x1, y1, x2, y2):
        # 找出两个单元格之间的墙并移除
        dx, dy = x2 - x1, y2 - y1
        
        if dx == 1:  # 右边
            self.grid[y1][x1]['walls'][3] = False  # 第一个单元格的右墙
            self.grid[y2][x2]['walls'][1] = False  # 第二个单元格的左墙
        elif dx == -1:  # 左边
            self.grid[y1][x1]['walls'][1] = False
            self.grid[y2][x2]['walls'][3] = False
        elif dy == 1:  # 下边
            self.grid[y1][x1]['walls'][2] = False  # 第一个单元格的下墙
            self.grid[y2][x2]['walls'][0] = False  # 第二个单元格的上墙
        elif dy == -1:  # 上边
            self.grid[y1][x1]['walls'][0] = False
            self.grid[y2][x2]['walls'][2] = False
    
    def step(self):
        if self.finished:
            return True
            
        if not self.stack:
            # 选择随机起始点
            x, y = random.randint(0, self.size-1), random.randint(0, self.size-1)
            self.grid[y][x]['visited'] = True
            self.stack.append((x, y))
            self.current_cell = (x, y)
            return False
        
        x, y = self.stack[-1]
        neighbors = self.get_neighbors(x, y)
        
        if not neighbors:
            self.stack.pop()
            if not self.stack:
                self.finished = True
                return True
            self.current_cell = self.stack[-1] if self.stack else None
            return False
        
        nx, ny, dx, dy = random.choice(neighbors)
        self.remove_wall(x, y, nx, ny)
        self.grid[ny][nx]['visited'] = True
        self.stack.append((nx, ny))
        self.current_cell = (nx, ny)
        return False
    
    def draw(self, screen):
        screen.fill(WHITE)
        
        # 绘制网格线
        for y in range(self.size):
            for x in range(self.size):
                cell = self.grid[y][x]
                
                # 绘制单元格
                if cell['visited']:
                    pygame.draw.rect(screen, WHITE, (x*CELL_SIZE, y*CELL_SIZE, CELL_SIZE, CELL_SIZE))
                
                # 绘制墙壁
                if cell['walls'][0]:  # 上墙
                    pygame.draw.line(screen, BLACK, (x*CELL_SIZE, y*CELL_SIZE), 
                                    ((x+1)*CELL_SIZE, y*CELL_SIZE), 2)
                if cell['walls'][1]:  # 左墙
                    pygame.draw.line(screen, BLACK, (x*CELL_SIZE, y*CELL_SIZE), 
                                    (x*CELL_SIZE, (y+1)*CELL_SIZE), 2)
                if cell['walls'][2]:  # 下墙
                    pygame.draw.line(screen, BLACK, (x*CELL_SIZE, (y+1)*CELL_SIZE), 
                                    ((x+1)*CELL_SIZE, (y+1)*CELL_SIZE), 2)
                if cell['walls'][3]:  # 右墙
                    pygame.draw.line(screen, BLACK, ((x+1)*CELL_SIZE, y*CELL_SIZE), 
                                    ((x+1)*CELL_SIZE, (y+1)*CELL_SIZE), 2)
        
        # 绘制当前单元格
        if self.current_cell:
            x, y = self.current_cell
            pygame.draw.rect(screen, GREEN, (x*CELL_SIZE+2, y*CELL_SIZE+2, CELL_SIZE-4, CELL_SIZE-4))
        
        # 绘制栈中的单元格（路径）
        for i, (x, y) in enumerate(self.stack):
            hex_color = 0xffffff - i*2
            color = ((hex_color >> 16) & 0xff, (hex_color >> 8) & 0xff, hex_color & 0xff)  # 渐变的绿色
            pygame.draw.rect(screen, color, (x*CELL_SIZE+2, y*CELL_SIZE+2, CELL_SIZE-4, CELL_SIZE-4))

# 创建迷宫
maze = Maze(GRID_SIZE)

# 主循环
running = True
clock = pygame.time.Clock()
auto_step = False
step_delay = 0.1  # 秒

while running:
    for event in pygame.event.get():
        if event.type == pygame.QUIT:
            running = False
        elif event.type == pygame.KEYDOWN:
            if event.key == pygame.K_SPACE:
                maze.step()
            elif event.key == pygame.K_r:
                maze.reset()
            elif event.key == pygame.K_a:
                auto_step = not auto_step
    
    if auto_step:
        finished = maze.step()
        if finished:
            auto_step = False
        time.sleep(step_delay)
    
    maze.draw(screen)
    pygame.display.flip()
    clock.tick(60)

pygame.quit()
